#!/usr/bin/env python3
"""
把 xyz 文件里的常见键名批量换成 MACE 官方默认键名
（energy → REF_energy，forces → REF_forces，stress → REF_stress...）
输出到 ./mace_ready/ 文件夹，文件名不变。
"""
import re
from pathlib import Path

# 旧 → 新
key_map = {
    'energy': 'REF_energy',
    'forces': 'REF_forces',
    'stress': 'REF_stress',
    'virials': 'REF_virials',
    'dipole': 'dipole',
    'polarizability': 'polarizability',
    'charges': 'REF_charges'
}


def process_properties(properties_str):
    """处理Properties字段，替换其中的键名"""
    parts = properties_str.split(':')
    result_parts = []

    i = 0
    while i < len(parts):
        # 键名在奇数位置 (0, 2, 4, ...)
        if i % 2 == 0 and i < len(parts):
            key = parts[i]
            # 如果是第一个部分"Properties"，不处理
            if i == 0 and key == "Properties":
                result_parts.append(key)
            else:
                # 替换键名
                new_key = key_map.get(key, key)
                result_parts.append(new_key)
        else:
            # 类型和维度信息直接保留
            result_parts.append(parts[i])
        i += 1

    return ':'.join(result_parts)


def repl_line(line: str) -> str:
    """把一行里的旧键全部替换"""
    # 处理Properties字段
    if 'Properties=' in line:
        # 使用正则匹配Properties=后面的整个值
        props_match = re.search(r'Properties=([^ ]+)', line)
        if props_match:
            old_props = props_match.group(1)
            new_props = process_properties(old_props)
            line = line.replace(f'Properties={old_props}', f'Properties={new_props}')

    # 处理其他键值对
    for old_key, new_key in key_map.items():
        # 匹配 key=value 或 key="value" 格式
        pattern = r'\b' + re.escape(old_key) + r'\s*=\s*("[^"]*"|[^\s]+)'
        line = re.sub(pattern, f'{new_key}=\\1', line)

    return line


# 主流程
if __name__ == '__main__':
    src_dir = Path('../mace_data_init')
    dst_dir = Path('../mace_data_ready')
    dst_dir.mkdir(exist_ok=True)

    xyz_files = list(src_dir.glob('*.xyz'))
    if not xyz_files:
        exit('当前目录下没有发现 *.xyz 文件')

    for f_in in xyz_files:
        f_out = dst_dir / f_in.name
        print(f'Processing {f_in}  ->  {f_out}')

        with f_in.open() as fin, f_out.open('w') as fout:
            for line in fin:
                # 替换帧头行（含 Lattice 或 Properties 的行）
                if line.strip() and ('Lattice' in line or 'Properties' in line):
                    line = repl_line(line)
                fout.write(line)

        # 验证替换结果
        with f_out.open('r') as f:
            content = f.read()
            if 'forces:' in content and 'REF_forces:' not in content:
                print(f"警告: {f_out} 中的 forces 可能未被正确替换")
            else:
                print(f"✓ {f_out} 处理完成")

    print(f'全部完成！结果在 {dst_dir.absolute()}')